//
// Created by ld016 on 16-12-27.
//

#ifndef VULKAN_LOADER_H
#define VULKAN_LOADER_H


namespace Vulkan{
    namespace Loader{
        Model loadModel(const std::string path) {
			std::vector<Vertex> vertices;
			std::vector<uint32_t> indices;
			Model model;
            auto processMesh = [&](aiMesh* mesh, const aiScene* scene)->Mesh
            {
                Mesh result;
                uint32_t indexBase = (uint32_t)indices.size();
                for (uint32_t i = 0; i < mesh->mNumVertices; i++)
                {
                    Vertex vertex;
                    vertex.position.x = mesh->mVertices[i].x;
                    vertex.position.y = mesh->mVertices[i].y;
                    vertex.position.z = mesh->mVertices[i].z;
                    // Normals
                    if (mesh->mNormals) {
                        vertex.normal.x = mesh->mNormals[i].x;
                        vertex.normal.y = mesh->mNormals[i].y;
                        vertex.normal.z = mesh->mNormals[i].z;
                    }
                    // Texture Coordinates
                    if (mesh->mTextureCoords[0])
                    {
                        vertex.uv.x = mesh->mTextureCoords[0][i].x;
                        vertex.uv.y = mesh->mTextureCoords[0][i].y;
                    }
                    vertices.push_back(vertex);
                    result.boundingBox.add(vertex.position);
                }
				model.boundingBox.add(result.boundingBox);

                for (uint32_t i = 0; i < mesh->mNumFaces; i++)
                {
                    aiFace face = mesh->mFaces[i];
                    for (uint32_t j = 0; j < face.mNumIndices; j++) {
                        uint32_t index = face.mIndices[j] + indexBase;
                        indices.push_back(index);
                    }
                }
                result.indexOffset = indexBase;
                result.indexCount = (uint32_t)indices.size()-indexBase;
                return result;
            };

            std::function<void(aiNode*,const aiScene*)> processNode = [&](aiNode* node, const aiScene* scene) {
                for (uint32_t i = 0; i < node->mNumMeshes; i++)
                {
                    aiMesh* aimesh = scene->mMeshes[node->mMeshes[i]];
					auto mesh = processMesh(aimesh, scene);
					memcpy(&mesh.transform, &node->mTransformation, sizeof(mesh.transform));
                    model.meshes.push_back(mesh);
                }
                for (uint32_t i = 0; i < node->mNumChildren; i++)
                {
                    processNode(node->mChildren[i], scene);
                }
            };

            Assimp::Importer import;
            const aiScene* scene = import.ReadFile(path, aiProcess_Triangulate | aiProcess_FlipUVs | aiProcess_GenNormals);

            if (!scene || scene->mFlags == AI_SCENE_FLAGS_INCOMPLETE || !scene->mRootNode)
            {
                throw std::runtime_error("ASSIMP ERROR : load model");
            }
            processNode(scene->mRootNode, scene);
			uint32_t vertexSize = (uint32_t)(vertices.size() * sizeof(Vertex));
			uint32_t indexSize = (uint32_t)(indices.size() * sizeof(uint32_t));
			createBufferMemoryWithStaging(vertexSize, vertices.data(), model.vertex.buffer, model.vertex.memory, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT);
			createBufferMemoryWithStaging(indexSize, indices.data(), model.index.buffer, model.index.memory, VK_BUFFER_USAGE_INDEX_BUFFER_BIT);
			model.indexCount = (uint32_t)indices.size();
            return model;
        }

        inline VkSampler createSample(float maxLod = 0.0f) {
            // Create sampler
            VkSamplerCreateInfo info = {};
            info.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
            info.magFilter = VK_FILTER_LINEAR;
            info.minFilter = VK_FILTER_LINEAR;
            info.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
            info.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
            info.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
            info.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
            info.mipLodBias = 0.0f;
            info.compareOp = VK_COMPARE_OP_NEVER;
            info.minLod = 0.0f;
            // Max level-of-detail should match mip level count
            info.maxLod = maxLod;
            // Enable anisotropic filtering
            info.maxAnisotropy = 8;
            info.anisotropyEnable = VK_TRUE;
            info.borderColor = VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE;

            VkSampler sampler = VK_NULL_HANDLE;
            vkCreateSampler(context.device, &info, nullptr, &sampler);
            return sampler;
        }

        inline VkImageView createImageView(VkImage& image, VkFormat format, 
			uint32_t mipLevel,uint32_t layercount,
			VkImageViewType type = VK_IMAGE_VIEW_TYPE_2D
		) {
            VkImageViewCreateInfo info = {};
            info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            info.viewType = type;
            info.format = format;
            info.components = { VK_COMPONENT_SWIZZLE_R, VK_COMPONENT_SWIZZLE_G, VK_COMPONENT_SWIZZLE_B, VK_COMPONENT_SWIZZLE_A };
            info.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, mipLevel, 0, layercount };
            info.image = image;
            VkImageView view = VK_NULL_HANDLE;
            vkCreateImageView(context.device, &info, nullptr, &view);
            return view;
        }

		VkImage createImage(VkExtent3D extent,uint32_t level,uint32_t layer,VkFormat format) {
			VkImageCreateInfo imageCreateInfo = {};
			imageCreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
			imageCreateInfo.imageType = VK_IMAGE_TYPE_2D;
			imageCreateInfo.format = format;
			imageCreateInfo.mipLevels = level;
			imageCreateInfo.arrayLayers = layer;
			imageCreateInfo.samples = VK_SAMPLE_COUNT_1_BIT;
			imageCreateInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
			imageCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
			imageCreateInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
			imageCreateInfo.extent = extent;
			imageCreateInfo.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;
			VkImage image = VK_NULL_HANDLE;
			vkCreateImage(context.device, &imageCreateInfo, nullptr, &image);
			return image;
		}

        std::vector<VkBufferImageCopy> imageCopy(gli::texture2d tex){
            std::vector<VkBufferImageCopy> bufferCopyRegions;
            uint32_t offset = 0;
            for (int i = 0; i < tex.levels(); ++i) {
                auto extent = tex.extent(i);
                VkBufferImageCopy copy = {};
                copy.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT,(uint32_t)i,(uint32_t)0,1};
                copy.imageExtent = {(uint32_t)extent.x,(uint32_t)extent.y,1};
                copy.bufferOffset = offset;
                offset += (uint32_t)tex.size(i);
                bufferCopyRegions.push_back(copy);
            }
            return bufferCopyRegions;
        }

        std::vector<VkBufferImageCopy> imageCopy(gli::texture_cube tex){
            std::vector<VkBufferImageCopy> bufferCopyRegions;
            uint32_t offset = 0;
            for (uint32_t face = 0; face < 6; face++)
            {
                for (uint32_t level = 0; level < tex.levels(); level++)
                {
                    auto extent = tex[face][level].extent();
                    VkBufferImageCopy copy = {};
                    copy.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT,(uint32_t)level,(uint32_t)face,1 };
                    copy.imageExtent = {(uint32_t)extent.x,(uint32_t)extent.y,1};
                    copy.bufferOffset = offset;
                    offset += (uint32_t)tex.size(level);
                    bufferCopyRegions.push_back(copy);
                }
            }
            return bufferCopyRegions;
        }

        std::vector<VkBufferImageCopy> imageCopy(gli::texture2d_array tex){
            std::vector<VkBufferImageCopy> bufferCopyRegions;
            bool sameDims = true;
            for (uint32_t layer = 0; layer < tex.layers(); layer++)
            {
                if (tex[layer].extent().x != tex.extent().x || tex[layer].extent().y != tex.extent().y)
                {
                    sameDims = false;
                    break;
                }
            }
            // If all layers of the texture array have the same dimensions, we only need to do one copy
            if (sameDims)
            {
                VkBufferImageCopy bufferCopyRegion = {};
                bufferCopyRegion.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT,0,0,(uint32_t)tex.layers()};
                bufferCopyRegion.imageExtent = {(uint32_t)tex.extent().x,(uint32_t)tex.extent().y,1};
                bufferCopyRegion.bufferOffset = 0;
                bufferCopyRegions.push_back(bufferCopyRegion);
            }
            else
            {
                uint32_t offset = 0;
                // If dimensions differ, copy layer by layer and pass offsets
                for (uint32_t layer = 0; layer < tex.layers(); layer++)
                {
                    VkBufferImageCopy bufferCopyRegion = {};
                    bufferCopyRegion.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT,0,layer,1};
                    bufferCopyRegion.imageExtent = {(uint32_t)tex.extent(layer).x,(uint32_t)tex.extent(layer).y,1};
                    bufferCopyRegion.bufferOffset = offset;
                    bufferCopyRegions.push_back(bufferCopyRegion);

                    offset += (uint32_t)tex.size(layer);
                }
            }
            return bufferCopyRegions;
        }

        Texture loadTexture(std::string path, VkFormat format = VK_FORMAT_R8G8B8A8_UNORM)
        {
            auto tex = gli::load(path.c_str());
            uint32_t width = tex.extent(0).x, height = tex.extent(0).y;

            VkBuffer stagingBuffer = createBuffer(tex.size(),VK_BUFFER_USAGE_TRANSFER_SRC_BIT);;
            VkDeviceMemory stagingMemory = createMemory(stagingBuffer,VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
            mapData(stagingMemory,tex.size(),tex.data());

			VkImageViewType viewType;
			uint32_t level = (uint32_t)tex.levels();
			uint32_t layer = (uint32_t)tex.layers();

			std::vector<VkBufferImageCopy> bufferCopyRegions;
            switch(tex.target()){
                case gli::TARGET_2D:
					bufferCopyRegions = imageCopy(gli::texture2d(tex));
					viewType = VK_IMAGE_VIEW_TYPE_2D;
					layer = 1;
                    break;
                case gli::TARGET_2D_ARRAY:
					bufferCopyRegions = imageCopy(gli::texture2d_array(tex));
					viewType = VK_IMAGE_VIEW_TYPE_2D_ARRAY;
                    break;
                case gli::TARGET_CUBE:
					bufferCopyRegions = imageCopy(gli::texture_cube(tex));
					viewType = VK_IMAGE_VIEW_TYPE_CUBE;
					layer = 6;
                    break;
				default:
					throw std::runtime_error("error image type");
					break;
            }

            // Create optimal tiled target image
			VkImage image = createImage({width,height,1},level,layer, format);
            VkDeviceMemory memory = createMemory(image,VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            VkImageSubresourceRange range = {VK_IMAGE_ASPECT_COLOR_BIT,0,level,0,layer};

            auto cmd = beginSingleCommandBuffer();
            setImageLayout(
                    cmd,
                    image,
                    VK_IMAGE_ASPECT_COLOR_BIT,
                    VK_IMAGE_LAYOUT_UNDEFINED,
                    VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                    range);

            // Copy mip levels from staging buffer
            vkCmdCopyBufferToImage(
                    cmd,
                    stagingBuffer,
                    image,
                    VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                    static_cast<uint32_t>(bufferCopyRegions.size()),
                    bufferCopyRegions.data()
            );

            setImageLayout(
                    cmd,
                    image,
                    VK_IMAGE_ASPECT_COLOR_BIT,
                    VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                    VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                    range);

            // Submit command buffer containing copy and image layout commands
            flushSingleCommandBuffer(cmd);
            // Clean up staging resources
            vkFreeMemory(context.device, stagingMemory, nullptr);
            vkDestroyBuffer(context.device, stagingBuffer, nullptr);

            auto sampler = createSample((float)level);
            auto view = createImageView(image, format, level, layer,viewType);

            VkDescriptorImageInfo descriptor = {sampler,view,VK_IMAGE_LAYOUT_GENERAL};

            return {image,memory,descriptor};
        }
    }
}


#endif
